"""
无人机避碰博弈问题（3D动力学模型）

两架无人机需要安全通过彼此，经过各自的途径点
- 无人机1：从西边(-5, 0, 25)运动，经过途径点(5, 0, 35)
- 无人机2：从北边(0, 5, 35)运动，经过途径点(0, -5, 25)

注意：途径点不是终点！无人机不需要停止，只需要安全通过即可。

使用非线性动力学模型:
状态: [px, py, pz, vx, vy, vz] (位置和速度)
控制: [theta, phi, tau] (俯仰角、滚转角、推力)
"""

import torch
import numpy as np
import casadi as ca
from typing import Dict, Tuple, Optional

# 解决Qt平台插件问题 - 使用Agg后端（非交互式）
import matplotlib
matplotlib.use('Agg')


class DroneGame:
    """
    两架无人机的非合作博弈（3D动力学）
    
    状态: [px1, py1, pz1, vx1, vy1, vz1, px2, py2, pz2, vx2, vy2, vz2]
    控制: [theta1, phi1, tau1, theta2, phi2, tau2] (俯仰角、滚转角、推力)
    """
    
    def __init__(
        self,
        dt: float = 0.1,
        collision_radius: float = 1.0,
        control_limits: Dict[str, float] = None,
        velocity_limit: float = 15.0,
        altitude_limits: Tuple[float, float] = None,
        g: float = 9.81
    ):
        """
        初始化无人机博弈
        
        Args:
            dt: 时间步长
            collision_radius: 碰撞半径
            control_limits: 控制输入限制 {'theta': 0.5, 'phi': 0.5, 'tau': 15.0}
            velocity_limit: 速度限制（默认15.0 m/s）
            altitude_limits: 高度限制 (min_altitude, max_altitude)，默认(15.0, 55.0)
            g: 重力加速度
        """
        self.n_players = 2
        self.n_states = 12  # 每个无人机6个状态(px, py, pz, vx, vy, vz)
        self.n_controls = 6  # 总控制维度：两个无人机各3个控制(theta, phi, tau)
        self.state_dim_per_player = 6  # 每个玩家6个状态
        self.control_dim_per_player = 3  # 每个玩家3个控制
        
        self.dt = dt
        self.collision_radius = collision_radius
        self.velocity_limit = velocity_limit
        self.g = g
        
        # 高度限制
        if altitude_limits is None:
            self.altitude_limits = (20.0, 40.0)  # 默认：最低20米，最高40米
        else:
            self.altitude_limits = altitude_limits
        
        # 控制限制
        if control_limits is None:
            self.control_limits = {
                'theta': 0.5,   # 俯仰角限制 (弧度)
                'phi': 0.5,     # 滚转角限制 (弧度)
                'tau': 25.0     # 推力限制 (m/s^2) - 增加以避免Z方向饱和
            }
        else:
            self.control_limits = control_limits
        
        # 目标位置 (3D) - 实际上是途径点，不是终点
        # 无人机需要经过这些点，但不需要停止
        self.target_positions = torch.tensor([
            [7.0, 0.0, 35.0],     # 无人机1途径点：东边5m，高度35米
            [0.0, -7.0, 25.0]    # 无人机2途径点：南边5m，高度25米
        ])
        
        # 初始位置 (3D) - 相近高度增加Z方向碰撞风险，使Z目标可观测
        self.initial_positions = torch.tensor([
            [-7.0, 0.0, 25.0],    # 无人机1起点：西边7m，高度30米
            [0.0, 7.0, 35.0]      # 无人机2起点：北边7m，高度30米（同高度）
        ])
        
        # 创建动力学模型
        self._create_dynamics()
    
    def _create_dynamics(self):
        """创建无人机动力学模型（使用CasADi）"""
        # 定义符号变量
        x = ca.MX.sym('x', 6)  # 单个无人机状态: [px, py, pz, vx, vy, vz]
        u = ca.MX.sym('u', 3)  # 控制: [theta, phi, tau]
        
        # 连续动力学方程
        xdot = ca.vertcat(
            x[3],                           # px_dot = vx
            x[4],                           # py_dot = vy
            x[5],                           # pz_dot = vz
            self.g * ca.tan(u[0]),         # vx_dot = g*tan(theta)
            -self.g * ca.tan(u[1]),        # vy_dot = -g*tan(phi)
            u[2] - self.g                  # vz_dot = tau - g
        )
        
        # 创建连续动力学函数
        self.dynamics_func_casadi = ca.Function('dynamics', 
                                                [x, u], [xdot], 
                                                ['x', 'u'], ['xdot'])
        
        # 离散化（RK4方法）
        self._discretize_dynamics(x, u, xdot)
    
    def _discretize_dynamics(self, x, u, xdot_expr):
        """使用RK4方法离散化动力学"""
        dt = self.dt
        
        # RK4积分
        k1 = xdot_expr
        k2 = ca.substitute(xdot_expr, x, x + dt/2 * k1)
        k3 = ca.substitute(xdot_expr, x, x + dt/2 * k2)
        k4 = ca.substitute(xdot_expr, x, x + dt * k3)
        
        x_next = x + dt/6 * (k1 + 2*k2 + 2*k3 + k4)
        
        # 创建离散步进函数
        self.step_func_casadi = ca.Function('step', 
                                            [x, u], [x_next], 
                                            ['x', 'u'], ['x_next'])
    
    def get_initial_state(self) -> torch.Tensor:
        """返回初始状态"""
        state = torch.zeros(self.n_states)
        # 无人机1初始状态 [px, py, pz, vx, vy, vz]
        state[0:3] = self.initial_positions[0]  
        state[3:6] = 0.0  # 初始速度为0
        # 无人机2初始状态
        state[6:9] = self.initial_positions[1]
        state[9:12] = 0.0
        return state
    
    def step_func(
        self,
        state: torch.Tensor,
        controls: torch.Tensor
    ) -> torch.Tensor:
        """
        动力学步进函数（纯PyTorch版本，保持梯度）
        
        🔧 关键修复：使用纯PyTorch实现，不调用.detach()，保持梯度流
        
        Args:
            state: 当前状态 [12] - [px1, py1, pz1, vx1, vy1, vz1, px2, py2, pz2, vx2, vy2, vz2]
            controls: 控制输入 [6] - [theta1, phi1, tau1, theta2, phi2, tau2]
            
        Returns:
            下一状态 [12]
        """
        # 🔧 使用纯PyTorch实现，保持可微分性
        return self._step_func_pytorch(state, controls)
    
    def _step_func_pytorch(
        self,
        state: torch.Tensor,
        controls: torch.Tensor
    ) -> torch.Tensor:
        """
        纯PyTorch实现的动力学步进（保持梯度流）
        
        动力学方程（欧拉法离散化）：
        px_dot = vx
        py_dot = vy
        pz_dot = vz
        vx_dot = g*tan(theta)
        vy_dot = -g*tan(phi)
        vz_dot = tau - g
        
        Args:
            state: [12] - [px1, py1, pz1, vx1, vy1, vz1, px2, py2, pz2, vx2, vy2, vz2]
            controls: [6] - [theta1, phi1, tau1, theta2, phi2, tau2]
            
        Returns:
            next_state: [12]
        """
        # 提取无人机1的状态和控制
        x1 = state[0:6]  # [px, py, pz, vx, vy, vz]
        u1 = controls[0:3]  # [theta, phi, tau]
        
        # 提取无人机2的状态和控制
        x2 = state[6:12]
        u2 = controls[3:6]
        
        # 无人机1的动力学
        x1_next = self._single_drone_dynamics_pytorch(x1, u1)
        
        # 无人机2的动力学
        x2_next = self._single_drone_dynamics_pytorch(x2, u2)
        
        # 组合成完整状态
        next_state = torch.cat([x1_next, x2_next])
        
        return next_state
    
    def _single_drone_dynamics_pytorch(
        self,
        x: torch.Tensor,
        u: torch.Tensor
    ) -> torch.Tensor:
        """
        单个无人机的动力学（纯PyTorch，欧拉法）
        
        Args:
            x: [6] - [px, py, pz, vx, vy, vz]
            u: [3] - [theta, phi, tau]
            
        Returns:
            x_next: [6]
        """
        # 提取状态
        px, py, pz = x[0], x[1], x[2]
        vx, vy, vz = x[3], x[4], x[5]
        
        # 提取控制
        theta, phi, tau = u[0], u[1], u[2]
        
        # 动力学方程
        px_dot = vx
        py_dot = vy
        pz_dot = vz
        vx_dot = self.g * torch.tan(theta)
        vy_dot = -self.g * torch.tan(phi)
        vz_dot = tau - self.g
        
        # 欧拉法离散化：x_next = x + dt * x_dot
        px_next = px + self.dt * px_dot
        py_next = py + self.dt * py_dot
        pz_next = pz + self.dt * pz_dot
        vx_next = vx + self.dt * vx_dot
        vy_next = vy + self.dt * vy_dot
        vz_next = vz + self.dt * vz_dot
        
        # 组合成下一状态
        x_next = torch.stack([px_next, py_next, pz_next, vx_next, vy_next, vz_next])
        
        return x_next
    
    def get_control_bounds(self) -> Tuple[torch.Tensor, torch.Tensor]:
        """返回控制输入的边界"""
        lb = torch.tensor([
            -self.control_limits['theta'],
            -self.control_limits['phi'],
            self.g * 0.5,  # 推力下界（至少能抵消部分重力）
            -self.control_limits['theta'],
            -self.control_limits['phi'],
            self.g * 0.5  # 推力下界
        ])
        
        ub = torch.tensor([
            self.control_limits['theta'],
            self.control_limits['phi'],
            self.control_limits['tau'],  # 推力上界
            self.control_limits['theta'],
            self.control_limits['phi'],
            self.control_limits['tau']
        ])
        
        return lb, ub
    
    def visualize_trajectory(
        self,
        states: torch.Tensor,
        save_path: Optional[str] = None
    ):
        """
        可视化3D轨迹
        
        Args:
            states: 状态序列 [horizon+1, n_states]
            save_path: 保存路径
        """
        import matplotlib.pyplot as plt
        from mpl_toolkits.mplot3d import Axes3D
        
        states_np = states.detach().cpu().numpy()
        
        fig = plt.figure(figsize=(15, 5))
        
        # 3D轨迹图
        ax1 = fig.add_subplot(131, projection='3d')
        
        # 无人机1轨迹（西到东）
        ax1.plot(states_np[:, 0], states_np[:, 1], states_np[:, 2], 
                'b-', linewidth=2, label='Drone 1')
        ax1.scatter(states_np[0, 0], states_np[0, 1], states_np[0, 2], 
                   c='blue', marker='o', s=100, label='Drone 1 Start')
        ax1.scatter(states_np[-1, 0], states_np[-1, 1], states_np[-1, 2], 
                   c='blue', marker='s', s=100, label='Drone 1 End')
        
        # 无人机2轨迹（北到南）
        ax1.plot(states_np[:, 6], states_np[:, 7], states_np[:, 8], 
                'r-', linewidth=2, label='Drone 2')
        ax1.scatter(states_np[0, 6], states_np[0, 7], states_np[0, 8], 
                   c='red', marker='o', s=100, label='Drone 2 Start')
        ax1.scatter(states_np[-1, 6], states_np[-1, 7], states_np[-1, 8], 
                   c='red', marker='s', s=100, label='Drone 2 End')
        
        # 目标位置
        ax1.scatter(self.target_positions[0, 0], self.target_positions[0, 1], 
                   self.target_positions[0, 2], c='blue', marker='*', s=300, 
                   label='Drone 1 Target')
        ax1.scatter(self.target_positions[1, 0], self.target_positions[1, 1], 
                   self.target_positions[1, 2], c='red', marker='*', s=300, 
                   label='Drone 2 Target')
        
        ax1.set_xlabel('X (m)')
        ax1.set_ylabel('Y (m)')
        ax1.set_zlabel('Z (m)')
        ax1.set_title('3D Trajectory')
        ax1.legend()
        ax1.grid(True)
        
        # XY平面投影
        ax2 = fig.add_subplot(132)
        ax2.plot(states_np[:, 0], states_np[:, 1], 'b-', linewidth=2, label='Drone 1')
        ax2.plot(states_np[:, 6], states_np[:, 7], 'r-', linewidth=2, label='Drone 2')
        ax2.scatter(self.target_positions[0, 0], self.target_positions[0, 1], 
                   c='blue', marker='*', s=300)
        ax2.scatter(self.target_positions[1, 0], self.target_positions[1, 1], 
                   c='red', marker='*', s=300)
        
        # 绘制碰撞区域
        for i in range(0, len(states_np), 5):
            circle1 = plt.Circle((states_np[i, 0], states_np[i, 1]), 
                                self.collision_radius, color='b', alpha=0.1)
            circle2 = plt.Circle((states_np[i, 6], states_np[i, 7]), 
                                self.collision_radius, color='r', alpha=0.1)
            ax2.add_patch(circle1)
            ax2.add_patch(circle2)
        
        ax2.set_xlabel('X (m)')
        ax2.set_ylabel('Y (m)')
        ax2.set_title('XY Plane Projection')
        ax2.legend()
        ax2.grid(True)
        ax2.axis('equal')
        
        # 高度变化
        ax3 = fig.add_subplot(133)
        time_steps = np.arange(len(states_np)) * self.dt
        ax3.plot(time_steps, states_np[:, 2], 'b-', linewidth=2, label='Drone 1 Altitude')
        ax3.plot(time_steps, states_np[:, 8], 'r-', linewidth=2, label='Drone 2 Altitude')
        ax3.axhline(y=self.target_positions[0, 2], color='b', linestyle='--', alpha=0.5)
        ax3.axhline(y=self.target_positions[1, 2], color='r', linestyle='--', alpha=0.5)
        ax3.set_xlabel('Time (s)')
        ax3.set_ylabel('Altitude (m)')
        ax3.set_title('Altitude Profile')
        ax3.legend()
        ax3.grid(True)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches='tight')
            print(f"轨迹图已保存到: {save_path}")
        
        plt.close()  # 关闭图形以释放内存
